"""
编码器
"""
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
import torch.nn as nn
import config


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=len(config.chatbot_ws_input),
                                      embedding_dim=config.chatbot_embedding_dim,
                                      padding_idx=config.chatbot_ws_input.PAD)
        self.gru = nn.GRU(input_size=config.chatbot_embedding_dim,
                          num_layers=config.chatbot_encoder_num_layers,
                          hidden_size=config.chatbot_encoder_hidden_size,
                          batch_first=True)

    def forward(self, input, input_length):
        """
        :param input:
        :param input_length:
        :return:
        """

        embeded = self.embedding(input)  # embeded: [batch_size,mmax_len,embedding_dim]

        embeded = pack_padded_sequence(embeded, input_length, batch_first=True)  # 打包

        out, hidden = self.gru(embeded)

        # 解包
        out, out_length = pad_packed_sequence(out, batch_first=True, padding_value=config.chatbot_ws_input.PAD)

        return out, hidden
